# face_detrec.py   （覆盖原文件）
import cv2
import numpy as np
import argparse
import time
import queue
import threading

from base.detection import Detection
from base.recognition import Recognition

from gpiozero.pins.lgpio import LGPIOFactory
from gpiozero import Device, Buzzer
Device.pin_factory = LGPIOFactory(chip=0)

# 有源蜂鸣器低电平触发 → 接 GPIO 70（根据你的引脚图）
buzzer = Buzzer(70, active_high=False)   # active_high=False 表示低电平响

# ---------------- 新增：欢迎界面线程 ----------------
import sys
from PyQt5.QtCore import QThread, pyqtSignal, QTimer
from PyQt5.QtWidgets import QApplication, QWidget, QLabel
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QFont


class WelcomeWindow(QWidget):
    def __init__(self):
        super().__init__()
        # 1. 去掉系统边框，确保置顶
        self.setWindowFlags(
            Qt.FramelessWindowHint |
            Qt.WindowStaysOnTopHint |
            Qt.Tool
        )
        self.setAttribute(Qt.WA_TranslucentBackground)

        # 2. 关键：拿到屏幕尺寸并手动铺满
        desktop = QApplication.primaryScreen().availableGeometry()
        self.setGeometry(desktop)

        # 3. 超大字体标签
        self.label = QLabel(self)
        self.label.setGeometry(desktop)          # 同样铺满
        self.label.setAlignment(Qt.AlignCenter)
        self.label.setStyleSheet("""
            color:#ffffff;
            background-color:rgba(0,0,0,220);
            border-radius:0px;
        """)
        font = QFont("微软雅黑", 120, QFont.Bold)  # 再调大
        self.label.setFont(font)

        # 4. 定时退出
        self.timer = QTimer(self)
        self.timer.setSingleShot(True)
        self.timer.timeout.connect(self.close)

    def popup(self, name: str):
        # 确保每次都重新拿到屏幕尺寸（防止外接显示器变化）
        desktop = QApplication.primaryScreen().availableGeometry()
        self.setGeometry(desktop)
        self.label.setGeometry(desktop)
        self.label.setText(f"欢迎回家，\n{name}！")
        self.show()
        self.raise_()          # 再置顶一次
        self.activateWindow()  # 强制激活
        self.timer.start(3000)


class UiThread(QThread):
    sig_popup = pyqtSignal(str)

    def __init__(self):
        super().__init__()
        self.app = QApplication(sys.argv)
        self.win = WelcomeWindow()
        self.sig_popup.connect(self.win.popup)

    def run(self):
        self.app.exec_()

# --------------------------------------------------


def face_detection_and_recognition(frame_queue, result_queue):
    det_model_path = "onnx_model/yolov5n-face_320_cut.q.onnx"
    rec_model_path = "onnx_model/arcface_mobilefacenet_cut.q.onnx"
    faces_path = "faces"

    det = Detection(det_model_path)
    rec = Recognition(rec_model_path, faces_path)

    while True:
        try:
            frame = frame_queue.get()
            face_imgs, boxes = det.infer_face(frame)
            results = {}
            if face_imgs is not None:
                for i, (face_img, box) in enumerate(zip(face_imgs, boxes)):
                    face_vector = rec.infer(face_img)
                    face_name = None
                    max_similarity_score = 0.0
                    for key, value in rec.face_bank.items():
                        similarity_scores = face_vector @ value.T
                        if similarity_scores[0][0] > max_similarity_score and similarity_scores[0][0] > 0.6:
                            max_similarity_score = similarity_scores[0][0]
                            face_name = key
                            results[face_name] = box
                        else:
                            results[f"unknown_{i}"] = box

            if results:
                # 只保留最新结果
                while not result_queue.empty():
                    result_queue.get()
                result_queue.put(results)

        except queue.Empty:
            continue


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--use_single_thread', action='store_true', help='use single thread')
    args = parser.parse_args()

    try:
        if not args.use_single_thread:
            cap = cv2.VideoCapture(20)
            ui = UiThread()
            ui.start()

            frame_queue = queue.Queue()
            result_queue = queue.Queue()
            threading.Thread(target=face_detection_and_recognition,
                             args=(frame_queue, result_queue),
                             daemon=True).start()

            # 用于防止重复弹窗
            already_welcomed = set()

            while True:
                ret, frame = cap.read()
                if not ret:
                    break

                # 只保留最新帧
                while not frame_queue.empty():
                    frame_queue.get()
                frame_queue.put(frame)

                # 取识别结果
                if not result_queue.empty():
                    results = result_queue.get()
                    current_names = set()
                    unknown_triggered = False  # 本轮是否响过

                    for name, box in results.items():
                        cv2.rectangle(frame, (int(box[0]), int(box[1])),
                                      (int(box[2]), int(box[3])), (0, 255, 0), 2)
                        if not name.startswith("unknown"):
                            cv2.putText(frame, name, (int(box[0]), int(box[1] - 10)),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
                            current_names.add(name)
                        else:
                            cv2.putText(frame, "unknown", (int(box[0]), int(box[1] - 10)),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
                            if not unknown_triggered:  # 只响一次
                                buzzer.on()
                                threading.Timer(0.2, buzzer.off).start()
                                unknown_triggered = True

                    # 新识别到的人 → 弹窗
                    for name in current_names:
                        if name not in already_welcomed:
                            ui.sig_popup.emit(name)
                            already_welcomed.add(name)

                    # 清理已离开画面的人
                    already_welcomed.intersection_update(current_names)

                cv2.imshow("Real-time Face Recognition", frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        else:
            # 单线程模式（保持不变）
            det_model_path = "onnx_model/yolov5n-face_320_cut.q.onnx"
            rec_model_path = "onnx_model/arcface_mobilefacenet_cut.q.onnx"
            faces_path = "faces"

            det = Detection(det_model_path)
            rec = Recognition(rec_model_path, faces_path)

            cap = cv2.VideoCapture(20)
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                frame_cp = frame.copy()
                face_imgs, boxes = det.infer_face(frame_cp)
                results = {}
                if len(face_imgs):
                    for i, (face_img, box) in enumerate(zip(face_imgs, boxes)):
                        face_vector = rec.infer(face_img)
                        face_name = None
                        max_similarity_score = 0.0
                        for key, value in rec.face_bank.items():
                            similarity_scores = face_vector @ value.T
                            if similarity_scores[0][0] > max_similarity_score and similarity_scores[0][0] > 0.6:
                                max_similarity_score = similarity_scores[0][0]
                                face_name = key
                                results[face_name] = box
                            else:
                                results[f"unknown_{i}"] = box
                for name, box in results.items():
                    cv2.rectangle(frame, (int(box[0]), int(box[1])),
                                  (int(box[2]), int(box[3])), (0, 255, 0), 2)
                    label = name if not name.startswith("unknown") else "unknown"
                    color = (0, 255, 0) if not name.startswith("unknown") else (0, 0, 255)
                    cv2.putText(frame, label, (int(box[0]), int(box[1] - 10)),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
                cv2.imshow("frame", frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

    finally:
        cv2.destroyAllWindows()


if __name__ == '__main__':
    main()